# -*- coding: utf-8 -*-
"""
@Auth ： 挂科边缘
@File ：detect.py
@IDE ：PyCharm
@Motto:学习新思想，争做新青年
@Email ：179958974@qq.com
"""

from ultralytics import YOLO
import os


def handle_picture(img_path):
    # Load a model
    try:
        model = YOLO(model=r'F:\ultralytics\runs\train\exp9\weights\best.pt')
    except:
        model = YOLO(model='yolo11n.pt')

    # Perform prediction and get the results
    results = model.predict(
        source=img_path,
        save=True,  # Save the results to disk
        show=False,  # Do not display the results on screen
    )

    # 获取保存路径
    save_dir = results[0].save_dir  # 获取第一个结果的保存目录
    save_path = os.path.join(save_dir, os.listdir(save_dir)[0])  # 获取保存文件的完整路径

    # 遍历检测结果并统计目标数量
    counts = {}  # 用于存储每个类别的计数

    for result in results:
        boxes = result.boxes  # 获取边界框信息
        class_ids = boxes.cls  # 获取类别ID

        # 统计每个类别的数量
        for cls_id in class_ids:
            if cls_id.item() in counts:
                counts[cls_id.item()] += 1
            else:
                counts[cls_id.item()] = 1
    output_str = "统计结果 "
    # 将检测结果追加到 output_str
    for i, count in counts.items():
        output_str += f" 天牛数目: {count} 个"

    return save_path  # 返回保存路径

def count_num(img_path):
    # Load a model
    try:
        model = YOLO(model=r'F:\ultralytics\runs\train\exp9\weights\best.pt')
    except:
        model = YOLO(model='yolo11n.pt')

    # Perform prediction and get the results
    results = model.predict(
        source=img_path,
        save=True,  # Save the results to disk
        show=False,  # Do not display the results on screen
    )

    # 遍历检测结果并统计目标数量
    counts = {}  # 用于存储每个类别的计数

    for result in results:
        boxes = result.boxes  # 获取边界框信息
        class_ids = boxes.cls  # 获取类别ID

        # 统计每个类别的数量
        for cls_id in class_ids:
            if cls_id.item() in counts:
                counts[cls_id.item()] += 1
            else:
                counts[cls_id.item()] = 1
    output_str = "统计结果 "
    # 将检测结果追加到 output_str
    for i, count in counts.items():
        output_str += f" 天牛数目: {count} 个"

    return output_str  # 返回保存路径
# 测试函数
if __name__ == "__main__":
    img_path = r"F:\天牛\VOCdevkit\images\train\tian (29).jpg"  # 替换为实际图片路径
    output_str = "检测结果:"  # 初始化 output_str
    save_path = handle_picture(img_path)  # 调用函数并获取保存路径
    print(f"保存路径: {save_path}")  # 打印保存路径
    print(output_str)  # 打印检测结果


